#!/bin/bash

set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
BIN_DIR="${SCRIPT_DIR}/../../bin"

TEMP_DIR="xllm"
INCLUDE_DIR="${TEMP_DIR}/include"
LIB_DIR="${TEMP_DIR}/lib"

VERSION_FILE="${SCRIPT_DIR}/../../version.txt"
TAR_BASE_NAME="xllm"
LOCAL_INSTALL_DIR="/usr/local"
LOCAL_TARGET_DIR="${LOCAL_INSTALL_DIR}/xllm"

HEADERS=("${SCRIPT_DIR}/llm.h" "${SCRIPT_DIR}/macros.h" "${SCRIPT_DIR}/types.h")
SO_FILES=(
    "${SCRIPT_DIR}/../../build/lib.linux-aarch64-cpython-311/xllm/libxllm.so"
    "/usr/local/Ascend/ascend-toolkit/8.2.RC1/opp/vendors/xllm/op_api/lib/libcust_opapi.so"
)

error_exit() {
    echo -e "\033[31merror: $1\033[0m" >&2
    exit 1
}

cd_bin_dir() {
    if [ ! -d "${BIN_DIR}" ]; then
        mkdir -p "${BIN_DIR}" || error_exit "failed to create bin directory: ${BIN_DIR}"
    fi
    
    cd "${BIN_DIR}" || error_exit "failed to enter bin directory: ${BIN_DIR}"
}

read_version() {
    if [ ! -f "${VERSION_FILE}" ]; then
        error_exit "${VERSION_FILE} is not existed"
    fi
    
    VERSION=$(cat "${VERSION_FILE}" | tr -d '[:space:]')
    if [ -z "${VERSION}" ]; then
        error_exit "version content is empty"
    fi
    
    TAR_FILE="${TAR_BASE_NAME}_${VERSION}.tar.gz"
}

check_files() {
    for header in "${HEADERS[@]}"; do
        if [ ! -f "${header}" ]; then
            error_exit "${header} is not existed"
        fi
    done
    
    for so_file in "${SO_FILES[@]}"; do
        if [ ! -f "${so_file}" ]; then
            error_exit "${so_file} is not existed"
        fi
    done
}

create_dirs() {
    mkdir -p "${INCLUDE_DIR}" || error_exit "create include directory failed"
    mkdir -p "${LIB_DIR}" || error_exit "create lib directory failed"
}

copy_headers() {
    for header in "${HEADERS[@]}"; do
        cp -f "${header}" "${INCLUDE_DIR}/" || error_exit "copy ${header} failed"
    done
}

copy_so() {
    for so_file in "${SO_FILES[@]}"; do
        cp -f "${so_file}" "${LIB_DIR}/" || error_exit "copy ${so_file} failed"
    done
}

package_tar() {
    tar -czf "${TAR_FILE}" "${TEMP_DIR}" || error_exit "tar failed"
}

cleanup_temp() {
    rm -rf "${TEMP_DIR}" || error_exit "rm temp directory failed"
}

extract_to_local() {
    if [ ! -f "${TAR_FILE}" ]; then
        error_exit "${TAR_FILE} is not existed"
    fi
    
    if [ ! -d "${LOCAL_INSTALL_DIR}" ]; then
        error_exit "local install directory is not existed"
    fi
    
    if [ -d "${LOCAL_TARGET_DIR}" ]; then
        rm -rf "${LOCAL_TARGET_DIR}" || error_exit "rm old xllm directory failed"
    fi

    tar -xzf "${TAR_FILE}" -C "${LOCAL_INSTALL_DIR}" || error_exit "extract failed"
}

main() {
    cd_bin_dir
    read_version
    check_files
    create_dirs
    copy_headers
    copy_so
    package_tar
    cleanup_temp
    extract_to_local
    
    echo -e "install file: \033[33m${TAR_FILE}\033[0m"
    echo -e "install path: \033[33m/usr/local/${TEMP_DIR}\033[0m"
}

main